# -*- coding: utf-8 -*-

import asyncio
import os
import re

import websockets
import uuid
import json
from tencentcloud.common import credential
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.lke.v20231130 import lke_client, models
import ssl
import certifi

from kj_d1.base.session import get_request_id, get_session


streaming_throttle = 10  # 节流控制

ssl_context = ssl.create_default_context()
ssl_context.load_verify_locations(certifi.where())

pattern = r'\d+(.*)'


class DSReplyTypeException(Exception):
    pass

def api_token(uc):
    # 实例化一个认证对象，入参需要传入腾讯云账户 SecretId 和 SecretKey，此处还需注意密钥对的保密
    # 代码泄露可能会导致 SecretId 和 SecretKey 泄露，并威胁账号下所有资源的安全性。以下代码示例仅供参考，建议采用更安全的方式来使用密钥，请参见：https://cloud.tencent.com/document/product/1278/85305
    # 密钥可前往官网控制台 https://console.cloud.tencent.com/cam/capi 进行获取
    cred = credential.Credential(uc.secret_id, uc.secret_key)
    # 实例化一个http选项，可选的，没有特殊需求可以跳过
    httpProfile = HttpProfile()
    httpProfile.endpoint = "lke.tencentcloudapi.com"

    # 实例化一个client选项，可选的，没有特殊需求可以跳过
    clientProfile = ClientProfile()
    clientProfile.httpProfile = httpProfile
    # 实例化要请求产品的client对象,clientProfile是可选的
    client = lke_client.LkeClient(cred, uc.region, clientProfile)

    # 实例化一个请求对象,每个接口都会对应一个request对象
    req = models.GetWsTokenRequest()
    params = {
        "Type": 5,
        "BotAppKey": uc.bot_app_key,
        "VisitorBizId": '001'
    }
    req.from_json_string(json.dumps(params))
    # 返回的resp是一个GetWsTokenResponse的实例，与请求对象对应
    resp = client.GetWsToken(req)
    # 输出json格式的字符串回包
    #print(resp.to_json_string())
    return resp.Token


async def call_ws(payload, token):
    url = f"wss://wss.lke.cloud.tencent.com/v1/qbot/chat/conn/?EIO=4&transport=websocket"

    async with websockets.connect(url, ssl=ssl_context) as ws:
        response = await ws.recv()
        #print("连接建立：" + response)

        auth = {"token": token}
        auth_message = f"40{json.dumps(auth)}"
        await ws.send(auth_message)

        # 收到服务端鉴权后的包
        response = await ws.recv()
        #print("鉴权结果：" + str(response))

        req_data = ["send", payload]

        send_data = f"42{json.dumps(req_data, ensure_ascii=False)}"
        #print("发送数据：" + send_data)
        await ws.send(send_data)
        is_kn = False
        while True:
            rsp = await ws.recv()
            if rsp == '2':
                # 收到了心跳包
                await ws.send("3")
                continue
            # 正则匹配 获取结果对象
            rsp_re_result = re.search(pattern, rsp).group(1)
            if '大模型直接回复' in rsp_re_result:
                raise DSReplyTypeException('不应为大模型直接回复，应为知识库调用')
            rsp_dict = json.loads(rsp_re_result)
            if rsp_dict[0] == "error":
                #print(rsp_dict)
                break
            elif rsp_dict[0] == "reply":
                if rsp_dict[1]["payload"]["is_from_self"]:
                    #print(f"from self:{rsp_dict[1]['payload']['content']}")
                    continue
                elif rsp_dict[1]["payload"]["is_final"]:
                    #print(f"数据接收完成::{rsp_dict[1]['payload']['content']}\n\n\n\n\n")
                    return rsp_dict[1]['payload']['content']
                else:
                    #print(rsp_dict)
                    continue


def ds_reply(uc, cont='', ):
    token = api_token(uc)
    if token == "":
        #print("get token error")
        exit(0)
    # "请分析该公司半年报中交易性金融资产、可供出售金融资产、持有至到期金融资产、外汇储备、外汇衍生品、境外债券投资相关的内容，整理出期末的详细情况。"
    payload = {
        "payload": {
            "request_id": get_request_id(),
            "session_id": get_session(),
            "content": cont
        }
    }

    raw1 = asyncio.run(call_ws(payload, token))
    return raw1


if __name__ == "__main__":
    from kj_d1.Config.config import useConfig
    rep = ds_reply(useConfig,
             "请参考万科A半年报中的投资状况、合并财务报表项目等内容，整理该公司的持有的外币货币性项目、投资外币衍生品的明细、持有而非发行的境外债券明细和QDII的投资情况，并总结成markdown表格，表格要求字段名是资产类别、期末价值(百万元)、持有明细。")
    print(rep)
